package tree.treerecursion;

import tree.TreeNode;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * @Classname : DeleteNodesAndReturnForest
 * @Description :
 * <a href="https://leetcode.cn/problems/delete-nodes-and-return-forest/">1110. 删点成林</a>
 * @Author : chentianyu
 * @Date 2023/1/4 22:05
 */


public class DeleteNodesAndReturnForest {
    Set<Integer> deletes;
    List<TreeNode> forest;

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        deletes = new HashSet<>();
        for (int d : to_delete) deletes.add(d);

        forest = new ArrayList<>();

        root = dfs(root);
        if (root != null) forest.add(root);
        return forest;
    }


    /**
     * 后序遍历实现
     *
     * @param cur
     * @return
     */
    private TreeNode dfs(TreeNode cur) {
        if (cur == null) return null;

        cur.left = dfs(cur.left);
        cur.right = dfs(cur.right);

        if (deletes.contains(cur.val)) {
            if (cur.left != null) forest.add(cur.left);
            if (cur.right != null) forest.add(cur.right);
            cur = null;
        }
        return cur;
    }
}
